import streamlit as st
import pandas as pd
import database as db
import cv2
import os
import numpy as np
import matplotlib.pyplot as plt
from analyzer import TrajectoryAnalyzer

def render():
    """渲染轨迹分析与可视化界面"""
    st.title("轨迹分析与可视化")
    
    # 检查是否选择了视频
    if 'selected_video_id' not in st.session_state:
        st.warning("请先从文件上传页面选择一个视频")
        if st.button("返回上传页面"):
            st.session_state.current_page = "upload"
            st.rerun()
        return
    
    # 获取视频信息
    video_id = st.session_state.selected_video_id
    
    # 获取轨迹数据
    trajectories = db.get_trajectories(video_id)
    
    if not trajectories:
        st.warning("该视频没有轨迹数据，请先处理视频")
        if st.button("返回检测页面"):
            st.session_state.current_page = "detection"
            st.rerun()
        return
    
    # 创建轨迹分析器
    analyzer = initialize_analyzer(trajectories, video_id)
    
    # 显示分析选项
    st.sidebar.header("分析选项")
    analysis_type = st.sidebar.selectbox(
        "选择分析类型",
        ["热力图", "轨迹统计", "轨迹图"]
    )
    
    # 显示不同的分析结果
    if analysis_type == "热力图":
        show_heatmap(analyzer)
    elif analysis_type == "轨迹统计":
        show_statistics(analyzer)
    elif analysis_type == "轨迹图":
        show_trajectory_plot(analyzer)
    
    # 添加返回按钮
    st.sidebar.markdown("---")
    if st.sidebar.button("返回视频检测页面"):
        st.session_state.current_page = "detection"
        st.rerun()

def initialize_analyzer(trajectories, video_id):
    """初始化轨迹分析器"""
    # 获取视频信息
    videos = db.get_videos()
    video_info = None
    for video in videos:
        if video['id'] == video_id:
            video_info = video
            break
    
    if not video_info:
        st.error("无法获取视频信息")
        return None
    
    # 获取视频尺寸
    video_path = video_info['processed_path'] or video_info['original_path']
    video_shape = get_video_shape(video_path)
    
    # 转换轨迹数据为DataFrame
    df = pd.DataFrame(trajectories)
    
    # 创建分析器
    analyzer = TrajectoryAnalyzer(df)
    if video_shape:
        analyzer.set_video_shape(video_shape[1], video_shape[0])  # width, height
    
    return analyzer

def get_video_shape(video_path):
    """获取视频尺寸"""
    if not os.path.exists(video_path):
        return None
    
    cap = cv2.VideoCapture(video_path)
    if not cap.isOpened():
        return None
    
    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
    cap.release()
    
    return (height, width)

def show_heatmap(analyzer):
    """显示热力图"""
    st.header("轨迹热力图")
    
    # 热力图参数
    col1, col2 = st.columns(2)
    with col1:
        resolution = st.slider("分辨率", min_value=50, max_value=200, value=100, step=10)
    
    with col2:
        sigma = st.slider("模糊度", min_value=1, max_value=10, value=2, step=1)
    
    # 生成热力图
    heatmap = analyzer.generate_heatmap(resolution=(resolution, resolution), sigma=sigma)
    
    # 显示热力图
    st.image(cv2.cvtColor(heatmap, cv2.COLOR_BGR2RGB), caption="行人轨迹热力图", use_column_width=True)
    
    # 显示密度图
    st.subheader("轨迹密度图")
    density_fig = analyzer.generate_density_map()
    st.pyplot(density_fig)

def show_statistics(analyzer):
    """显示统计数据"""
    st.header("轨迹统计数据")
    
    # 计算统计
    stats = analyzer.calculate_statistics()
    
    # 显示统计信息
    st.subheader("基本统计")
    st.write(f"- **轨迹总数**: {stats['total_tracks']}")
    st.write(f"- **轨迹点总数**: {stats['total_points']}")
    st.write(f"- **平均每轨迹点数**: {stats['avg_points_per_track']:.2f}")
    
    if 'avg_track_length' in stats:
        st.write(f"- **平均轨迹长度**: {stats['avg_track_length']:.2f} 像素")
        st.write(f"- **最大轨迹长度**: {stats['max_track_length']:.2f} 像素")
        st.write(f"- **最小轨迹长度**: {stats['min_track_length']:.2f} 像素")
    
    # 显示轨迹长度直方图
    st.subheader("轨迹长度分布")
    hist_fig = analyzer.plot_track_length_histogram()
    st.pyplot(hist_fig)

def show_trajectory_plot(analyzer):
    """显示轨迹图"""
    st.header("行人轨迹图")
    
    # 显示轨迹图
    fig = analyzer.plot_interactive_trajectories()
    st.pyplot(fig) 